# %% import 

import sys,copy,os,inspect
if hasattr(sys.modules[__name__], '__file__'):
    _file_name = __file__
else:
    _file_name = inspect.getfile(inspect.currentframe())
CURRENT_FILE_PATH = os.path.dirname(_file_name)
sys.path.append(os.getcwd()+"/../neuronVis")
import pandas as pd
import Scene
import BrainRegion as BR 
import IONData 
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
matplotlib.use('module://matplotlib_inline.backend_inline')
%matplotlib inline
iondata = IONData.IONData()

br = BR.BrainRegion()
br.praseJson()

# %% [markdown]
# # Get all the neuron info


d1samples=['230012','230013','233664','233665','240833','240834']
d2samples=['230015','233669','233670','241311','241312','241313']
ChATsamples=['230016','221459']
calcrsamples=['233675','233676']
tac2samples=['233381','233382']
sampleCre={'D1':d1samples,'D2':d2samples,'ChAT':ChATsamples,'CALCR':calcrsamples,'TAC2':tac2samples}

neuronCre={}

for cre,samples in sampleCre.items():
    neurons=[]
    for sample in samples:
        neuronOfSample =iondata.getNeuronListBySampleID(sample)
        somas=iondata.getSoma(sample)
        for neuron in neuronOfSample:
            for soma in somas:
                if soma['sampleid']==neuron['sampleid'] and soma['name']==neuron['name']:
                    neuron['xyz']=[soma['somax'],soma['somay'],soma['somaz']]
                    break
            
        neurons.extend(neuronOfSample)

    neuronCre[cre]=neurons

# %% [markdown]
# # create scene

import nrrd
res = iondata.getStructureMask(5601)
maskdata,header = nrrd.read(res[1],index_order='F')

neurons=[]
for cre,neuronsinCre in neuronCre.items():
    neurons+=neuronsinCre
    ACBneurons=[]
    Otherneurons=[]
    for neuron in neuronsinCre:
        if neuron['region']=='ACB':
            neuron['color']={'r':'255','g':'0','b':'0'}
            ACBneurons.append(neuron)
        else:
            neuron['color']={'r':'0','g':'200','b':'0'}
            Otherneurons.append(neuron)
        if maskdata[int(neuron['xyz'][0]/10),int(neuron['xyz'][1]/10),int(neuron['xyz'][2]/10)]:
            neuron['region']='ACBc'
            neuron['color']={'r':'0','g':'0','b':'200'}
    Scene.createScene(neuronlist=ACBneurons,filename='../resource/scene/STR'+cre+'-ACBneurons.nv')



# %% [markdown]
# # neurons in regions
regionNeurons={}
regions=[]
for neuron in neurons:
    if neuron['region'] not in regionNeurons.keys():
        regions.append(neuron['region'])
        regionNeurons[neuron['region']]=[]
    regionNeurons[neuron['region']].append(neuron)
count=[]
for region in regionNeurons.keys():
    count.append(len(regionNeurons[region]))
    Scene.createScene(neuronlist=regionNeurons[region],filename='../resource/scene/STR_'+region+'neurons.nv')
plt.bar(height=count,x=regions)
plt.title('Soma distribution')
plt.ylabel('Neuron Count')

plt.xticks(regions,regions, rotation=60)
plt.show()


# %% [markdown]
# # neuron distribution of cre
bottom = np.zeros(len(regions))
for cre,neuronsinCre in neuronCre.items():
    neuroninfodf = pd.DataFrame(neuronsinCre)
    neuronInRegion = neuroninfodf['region'].value_counts().to_dict()
    y=[neuronInRegion[region] if region in neuronInRegion.keys() else 0 for region in regions]
    plt.bar(x=range(0,len(regions)), height=y, bottom=bottom,label=cre)
    bottom+=y
plt.title('Soma distribution')
plt.ylabel('Neuron Count')
plt.legend()
plt.xticks(range(0,len(regions)),regions, rotation=60)
plt.show()
# %% 


# %% [markdown]
# # ACBneruons
ACBneurons=[]
for cre,neuronsinCre in neuronCre.items():
    if cre!='D1':
        continue
    for neuron in neuronsinCre:
        if neuron['region']=='ACB':
            neuron['color']={'r':'255','g':'0','b':'0'}
            ACBneurons.append(neuron)
acbdf = iondata.getPropertiesDF(ACBneurons)
acbprojectregiondf= acbdf[acbdf.property.str.contains('projectregion',case=False)]
acbprojectregiondf.set_index(['property'],inplace=True)
acbnewindex = acbprojectregiondf.sum(1).sort_values(ascending=False).index[0:55]
acbprojectregiondf=acbprojectregiondf.loc[acbnewindex,:]
acbdata =acbprojectregiondf.to_numpy().astype(np.float32)
acbdata=(np.log2(acbdata/100.0+1))

acbdf2 = pd.DataFrame(acbdata, index=acbprojectregiondf.index, columns=acbprojectregiondf.columns)
# fig,ax = plt.subplots(figsize=(13,7))
sns.clustermap(acbdf2,row_cluster=False,figsize=[15,25])


# %% [markdown]
# Other Neurons
# %% 

otherdf = iondata.getPropertiesDF(Otherneurons)
otherprojectregiondf= otherdf[otherdf.property.str.contains('projectregion',case=False)]
otherprojectregiondf.set_index(['property'],inplace=True)
othernewindex = otherprojectregiondf.sum(1).sort_values(ascending=False).index[0:55]
otherprojectregiondf=otherprojectregiondf.loc[othernewindex,:]
otherdata =otherprojectregiondf.to_numpy().astype(np.float32)
otherdata=(np.log2(otherdata/100.0+1))

otherdf2 = pd.DataFrame(otherdata, index=otherprojectregiondf.index, columns=otherprojectregiondf.columns)
sns.clustermap(otherdf2,figsize=[50,20])


# %% [markdown]
# # render single neuron

import Visual as nv

neuronvis = nv.neuronVis(size=(1700,1000),renderModel=0)

neuronvis.render.setBackgroundColor((1.0,1.0,1.,1.0))

neuronvis.setLineWidth(1)
neuronvis.addRegion('STR')
for neuron in Otherneurons:
    neuronvis.addNeuronByID(neuron['sampleid'],neuron['name'],color=[1,0,0],somaColor=[0,1,0],somaHide=False,axonHide=False,dendriteHide=False,isLine=True)
    neuronvis.render.savepng('../resource/png/ACB/STRother/'+neuron['region']+'/'+neuron['sampleid']+neuron['name']+'.png')
    neuronvis.clear(regions=False)
neuronvis.render.closeWindow()



# %% [markdown]
# # cluster all the acb neurons of D1
regionNeuronsD1={}
regions=[]

for neuron in neuronCre['D1']:
    if neuron['region'] not in regionNeuronsD1.keys():
        regions.append(neuron['region'])
        regionNeuronsD1[neuron['region']]=[]
    regionNeuronsD1[neuron['region']].append(neuron)

strNeurons=[]
strRegions=['ACB']
for region in strRegions:
    strNeurons.extend(regionNeuronsD1[region])

strdf = iondata.getPropertiesDF(strNeurons)
strprojectregiondf= strdf[strdf.property.str.contains('projectregion',case=False)]
strprojectregiondf.set_index(['property'],inplace=True)
strnewindex = strprojectregiondf.sum(1).sort_values(ascending=False).index[0:55]
strprojectregiondf=strprojectregiondf.loc[strnewindex,:]
strdata =strprojectregiondf.to_numpy().astype(np.float32)
strdata=(np.log2(strdata/100.0+1))

strdf2 = pd.DataFrame(strdata, index=strprojectregiondf.index, columns=strprojectregiondf.columns)

g=sns.clustermap(strdf2,figsize=[10,15] )
g.fig.suptitle('ACB neurons of D1') 

# %%
import scipy

def flood_fill_hull(image):    
    points = np.transpose(np.where(image))
    hull = scipy.spatial.ConvexHull(points)
    deln = scipy.spatial.Delaunay(points[hull.vertices]) 
    idx = np.stack(np.indices(image.shape), axis = -1)
    out_idx = np.nonzero(deln.find_simplex(idx) + 1)
    out_img = np.zeros(image.shape)
    out_img[out_idx] = 1
    return out_img, hull




points = tuple(np.rint(10 * np.random.randn(3,100)).astype(int) + 50)
image = np.zeros((100,)*3)
print(image.shape)
image[points] = 1


out, h = flood_fill_hull(image)

plt.imshow(out[50])
#%%

dendriteVolume={}

for key,neurons in regionNeuronsD1.items():
    if key!='ACBc' and key!='ACB':
        continue
    for neuron in neurons:
        tree =iondata.getNeuronTreeByID(neuron['sampleid'], neuron['name'])
        dendrite = tree.getDendrite()
        pointsx=[]
        pointsy=[]
        pointsz=[]
        for edge in dendrite:
            for p in edge.data:

                pointsx.append(p.x/10)
                pointsy.append(p.y/10)
                pointsz.append(p.z/10)
        if len(pointsx)==0:
            continue
        arrayx,arrayy,arrayz = np.array(pointsx),np.array(pointsy),np.array(pointsz)
        arrayx = arrayx-np.min(arrayx)
        arrayy = arrayy-np.min(arrayy)
        arrayz = arrayz-np.min(arrayz)
        points=tuple((np.rint(arrayx).astype(int),np.rint(arrayy).astype(int),np.rint(arrayz).astype(int)))
        image = np.zeros((int(arrayx.max()+2),int(arrayy.max()+2),int(arrayz.max()+2)))
        print(image.shape)
        image[points] = 1
        out, h = flood_fill_hull(image)
        dendriteVolume[neuron['sampleid']+neuron['name']]=np.sum(out)

# %% [markdown]
# # dendrite volume
import RenderGL,GeometryAdapter
import math
render = RenderGL.RenderGL(renderModel=0)
render.backGroundColor = [1.0,1.0,1.0]
iondata = IONData.IONData()

response,regionFileName=iondata.getFileFromServer('allobj/56.obj')
		# regionFileName = CURRENT_FILE_PATH+'/../resource/allobj/'+str(self.regionName2ID[name][0])+'.obj'
geo  = render.loadGeometry(regionFileName)
render.addGeometry(geo,[0,0,1])

response,regionFileName=iondata.getFileFromServer('allobj/5601.obj')
		# regionFileName = CURRENT_FILE_PATH+'/../resource/allobj/'+str(self.regionName2ID[name][0])+'.obj'
geo  = render.loadGeometry(regionFileName)
render.addGeometry(geo,[0,1,0])

	# render.setBackgroundColor([1,1,1])
	# render.setLineWidth(4.0)
ga = GeometryAdapter.GeometryAdapter()
index=0
maxvolume=0
minvolume=999
volumearray=[]
for neurons in [regionNeuronsD1[key] for key in['ACB','ACBc']]:
    for neuron in neurons:
        if neuron['sampleid']+neuron['name'] in dendriteVolume:
            volume = dendriteVolume[neuron['sampleid']+neuron['name']]
            point = [neuron['xyz'][0],neuron['xyz'][1],11400-neuron['xyz'][2] if neuron['xyz'][2]>5700 else neuron['xyz'][2]]
            if volume<20000:
                # if point[2]>5700:
                #     point[2]=11400-point[2]
                logvolume=np.log10(volume)
                volumearray.append(volume)
                r = (volume ** (1/6)-2.8)/(5-2.8)
                maxvolume=r if r>maxvolume else maxvolume
                minvolume=r if r<minvolume else minvolume

                print(neuron['sampleid']+neuron['name'],volume,logvolume,r)
                ga.geometry.addPoint(point,[r,r,r])

                ga.geometry.addIndex(index)
                index+=1

ga.geometry.drawModel='points'
render.setPointSize(12)
render.addGeometry(ga.geometry)
render.setView()

render.run()

# %%
